"""
RAG工具类

为Agent提供RAG相关功能的工具
"""

import os
import logging
from typing import Dict, List, Any, Optional
from langchain.tools import BaseTool
from pydantic import BaseModel, Field

from core.rag_manager import RAGManager

logger = logging.getLogger(__name__)

# 全局RAG管理器实例
_rag_manager = None

def get_rag_manager() -> RAGManager:
    """获取RAG管理器实例"""
    global _rag_manager
    if _rag_manager is None:
        _rag_manager = RAGManager()
    return _rag_manager

class DocumentQuestionInput(BaseModel):
    """文档问答输入"""
    question: str = Field(description="要询问的问题")

class DocumentSearchInput(BaseModel):
    """文档搜索输入"""
    query: str = Field(description="搜索查询")
    top_k: int = Field(default=5, description="返回结果数量")

class AddDocumentInput(BaseModel):
    """添加文档输入"""
    file_path: str = Field(description="文档文件路径")

class AddDocumentDirectoryInput(BaseModel):
    """添加文档目录输入"""
    directory_path: str = Field(description="文档目录路径")

class AddDocumentUrlsInput(BaseModel):
    """添加文档URL输入"""
    urls: List[str] = Field(description="URL列表")

class DocumentQuestionTool(BaseTool):
    """文档问答工具"""
    name = "document_question"
    description = "基于已加载的文档回答问题。适用于需要从文档中获取信息的查询。"
    args_schema = DocumentQuestionInput
    
    def _run(self, question: str) -> str:
        """执行文档问答"""
        try:
            rag_manager = get_rag_manager()
            result = rag_manager.ask_question(question)
            
            if result['success']:
                answer = result['answer']
                sources = result.get('sources', [])
                
                # 格式化回答
                formatted_answer = f"回答：{answer}\n"
                
                if sources:
                    formatted_answer += "\n参考来源：\n"
                    for i, source in enumerate(sources, 1):
                        formatted_answer += f"{i}. {source['source']}\n"
                        if source.get('content_preview'):
                            formatted_answer += f"   内容预览：{source['content_preview']}\n"
                
                return formatted_answer
            else:
                return f"无法回答问题：{result.get('answer', '未知错误')}"
                
        except Exception as e:
            logger.error(f"文档问答失败: {e}")
            return f"文档问答过程中发生错误：{e}"

class DocumentSearchTool(BaseTool):
    """文档搜索工具"""
    name = "document_search"
    description = "在已加载的文档中搜索相关内容。适用于查找特定信息或主题的文档片段。"
    args_schema = DocumentSearchInput
    
    def _run(self, query: str, top_k: int = 5) -> str:
        """执行文档搜索"""
        try:
            rag_manager = get_rag_manager()
            results = rag_manager.search_documents(query, k=top_k)
            
            if not results:
                return "未找到相关文档。请确保已加载文档到知识库中。"
            
            # 格式化搜索结果
            formatted_results = f"找到 {len(results)} 个相关文档片段：\n\n"
            
            for result in results:
                formatted_results += f"排名 {result['rank']}：\n"
                formatted_results += f"来源：{result['source']}\n"
                formatted_results += f"内容：{result['content'][:300]}{'...' if len(result['content']) > 300 else ''}\n"
                formatted_results += "-" * 50 + "\n\n"
            
            return formatted_results
            
        except Exception as e:
            logger.error(f"文档搜索失败: {e}")
            return f"文档搜索过程中发生错误：{e}"

class AddDocumentTool(BaseTool):
    """添加文档工具"""
    name = "add_document"
    description = "将单个文档文件添加到知识库中。支持PDF、Word、TXT、Markdown等格式。"
    args_schema = AddDocumentInput
    
    def _run(self, file_path: str) -> str:
        """添加文档文件"""
        try:
            # 检查文件是否存在
            if not os.path.exists(file_path):
                return f"文件不存在：{file_path}"
            
            rag_manager = get_rag_manager()
            result = rag_manager.add_documents_from_file(file_path)
            
            if result['success']:
                return f"成功添加文档：{result['message']}"
            else:
                return f"添加文档失败：{result['message']}"
                
        except Exception as e:
            logger.error(f"添加文档失败: {e}")
            return f"添加文档过程中发生错误：{e}"

class AddDocumentDirectoryTool(BaseTool):
    """添加文档目录工具"""
    name = "add_document_directory"
    description = "将目录中的所有支持格式的文档添加到知识库中。会递归搜索子目录。"
    args_schema = AddDocumentDirectoryInput
    
    def _run(self, directory_path: str) -> str:
        """添加文档目录"""
        try:
            # 检查目录是否存在
            if not os.path.exists(directory_path):
                return f"目录不存在：{directory_path}"
            
            if not os.path.isdir(directory_path):
                return f"路径不是目录：{directory_path}"
            
            rag_manager = get_rag_manager()
            result = rag_manager.add_documents_from_directory(directory_path)
            
            if result['success']:
                return f"成功添加目录文档：{result['message']}"
            else:
                return f"添加目录文档失败：{result['message']}"
                
        except Exception as e:
            logger.error(f"添加目录文档失败: {e}")
            return f"添加目录文档过程中发生错误：{e}"

class AddDocumentUrlsTool(BaseTool):
    """添加文档URL工具"""
    name = "add_document_urls"
    description = "从网页URL加载文档内容并添加到知识库中。支持多个URL。"
    args_schema = AddDocumentUrlsInput
    
    def _run(self, urls: List[str]) -> str:
        """添加文档URL"""
        try:
            if not urls:
                return "URL列表为空"
            
            rag_manager = get_rag_manager()
            result = rag_manager.add_documents_from_urls(urls)
            
            if result['success']:
                return f"成功添加URL文档：{result['message']}"
            else:
                return f"添加URL文档失败：{result['message']}"
                
        except Exception as e:
            logger.error(f"添加URL文档失败: {e}")
            return f"添加URL文档过程中发生错误：{e}"

class KnowledgeBaseInfoTool(BaseTool):
    """知识库信息工具"""
    name = "knowledge_base_info"
    description = "获取当前知识库的状态信息，包括文档数量、向量数量等统计信息。"
    
    def _run(self) -> str:
        """获取知识库信息"""
        try:
            rag_manager = get_rag_manager()
            info = rag_manager.get_knowledge_base_info()
            
            # 格式化信息
            formatted_info = "知识库状态信息：\n"
            formatted_info += f"- 是否已初始化：{'是' if info['has_vectorstore'] else '否'}\n"
            formatted_info += f"- 向量数量：{info['vector_count']}\n"
            formatted_info += f"- 文档数量：{info['statistics']['total_documents']}\n"
            formatted_info += f"- 文本块数量：{info['statistics']['total_chunks']}\n"
            formatted_info += f"- 最后更新：{info['statistics']['last_updated'] or '未知'}\n"
            
            if info['statistics']['sources']:
                formatted_info += "- 数据源：\n"
                for source in info['statistics']['sources']:
                    formatted_info += f"  * {source}\n"
            else:
                formatted_info += "- 数据源：无\n"
            
            return formatted_info
            
        except Exception as e:
            logger.error(f"获取知识库信息失败: {e}")
            return f"获取知识库信息过程中发生错误：{e}"

class ClearKnowledgeBaseTool(BaseTool):
    """清空知识库工具"""
    name = "clear_knowledge_base"
    description = "清空当前知识库中的所有文档和向量数据。此操作不可逆。"
    
    def _run(self) -> str:
        """清空知识库"""
        try:
            rag_manager = get_rag_manager()
            success = rag_manager.clear_knowledge_base()
            
            if success:
                return "知识库已成功清空"
            else:
                return "清空知识库失败"
                
        except Exception as e:
            logger.error(f"清空知识库失败: {e}")
            return f"清空知识库过程中发生错误：{e}"

class ExportKnowledgeBaseSummaryTool(BaseTool):
    """导出知识库摘要工具"""
    name = "export_knowledge_base_summary"
    description = "导出当前知识库的详细摘要报告，包括配置信息、统计数据等。"
    
    def _run(self) -> str:
        """导出知识库摘要"""
        try:
            rag_manager = get_rag_manager()
            summary = rag_manager.export_knowledge_base_summary()
            return summary
            
        except Exception as e:
            logger.error(f"导出知识库摘要失败: {e}")
            return f"导出知识库摘要过程中发生错误：{e}"

# 所有RAG工具列表
RAG_TOOLS = [
    DocumentQuestionTool(),
    DocumentSearchTool(),
    AddDocumentTool(),
    AddDocumentDirectoryTool(),
    AddDocumentUrlsTool(),
    KnowledgeBaseInfoTool(),
    ClearKnowledgeBaseTool(),
    ExportKnowledgeBaseSummaryTool(),
]

def get_rag_tools() -> List[BaseTool]:
    """获取所有RAG工具"""
    return RAG_TOOLS 